import sys
sys.path.append('..')
from Function.readdatafunction import readdata
from Function.CNN import SimpleCNN
from Function.loss import softmax
from Function.accuracy import comouteaccuracy
from Function.drawPicture import drawPicture
import torch

net = SimpleCNN(num_classes=16, in_channels=204)
PATH = './model/Salinas/CNN_Salinas.pt'
net.load_state_dict(torch.load(PATH))

Xtrain, Ytrain, X, Y, res = readdata("../dataset/Salinas/Salinas_corrected.mat", "salinas_corrected",
                                        "../dataset/Salinas/Salinas_gt.mat", "salinas_gt",
                                        17, 204, 16, "CNN")
del Xtrain, Ytrain

X = torch.from_numpy(X)
Y = torch.from_numpy(Y)
y_pred = softmax(net(X)).argmax(dim=1)

color = [[156, 134, 23], [0, 255, 50], [255, 190, 0], [0, 100, 255], [75, 23, 196],
         [63, 96, 236], [152, 16, 78], [255, 0, 36], [0, 223, 36], [76, 128, 214],
         [100, 62, 0], [0, 145, 186], [12, 190, 190], [22, 22, 22], [216, 20, 154],
         [213, 15, 42]]

comouteaccuracy(y_pred.numpy(), Y.numpy())
drawPicture(y_pred, res, 512, 217, './Picture/Salinas', 'CNN')

